import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt
import os
import sys
from python_ai.common.xcommon import *
from python_ai.CV_2.app.teachers_day17.my.inception_keras_oop_cifar10 import InceptionNet
import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import layers, activations, optimizers, losses
from sklearn.model_selection import train_test_split
import os

np.random.seed(3)
FILE_NAME = os.path.basename(__file__)

vcode_dir_train = '../../../../large_data/DL1/_many_files/vcode_data/train'
vcode_dir_test = '../../../../large_data/DL1/_many_files/vcode_data/test'

def load_dir(vcode_dir):
    names = os.listdir(vcode_dir)
    n_names = len(names)
    names = np.array(names)
    idx = np.random.permutation(n_names)  # for shuffle
    names = names[idx]
    x = []
    y = []
    for name in names:
        digits, ext = os.path.splitext(name)
        if ext.lower() != '.jpg':
            continue
        digits_list = list(digits)
        if len(digits_list) != 4:
            continue
        digits_arr = []
        try:
            for di in digits_list:
                digits_arr.append(int(di))
        except ValueError as ex:
            continue

        y.append(digits_arr)

        path = os.path.join(vcode_dir, name)
        img = cv.imread(path, cv.IMREAD_GRAYSCALE)
        x.append(img)

    x = np.float32(x) / 255.
    y = np.uint8(y)
    x = np.expand_dims(x, 3)
    y = keras.utils.to_categorical(y, 10)
    y = y.reshape((-1, 40))
    return x, y

x_train, y_train = load_dir(vcode_dir_train)
x_test, y_test = load_dir(vcode_dir_test)

model = InceptionNet(40, 16, 2, activation=None)
model.build(input_shape=(None, 60, 160, 1))
model.summary(line_length=120)

model.compile(loss=losses.BinaryCrossentropy(from_logits=True),
              optimizer=optimizers.Adam(lr=0.0001),
              metrics=['accuracy'])

history=model.fit(
    x_train, y_train, batch_size=64, epochs=20, validation_split=0.05,
    callbacks=[keras.callbacks.TensorBoard(log_dir='_log/' + FILE_NAME, update_freq='batch', profile_batch=0)],

)

score=model.evaluate(x_test, y_test, batch_size=64)
print('accuracy',score[1])
print('loss',score[0])


